# -*- coding: utf-8 -*-
import asyncio
from datetime import datetime
from threading import Thread
from typing import Dict
from fastapi import WebSocket

from apps.ws.schemas import Message
from application import settings
from core.redis import RedisManager
from core.logger import logger
from utils.tools import SingletonDecorator


class Connection:

  def __init__(
      self,
      client_id: str,
      websocket: WebSocket
  ):
    self.websocket = websocket
    self.client_id = client_id

  async def send_message(self, message: Message):
    await self.websocket.send_text(f"{message.model_dump_json()}")

  async def ping(self):
    await self.websocket.send_text(f"ping {settings.server_id} {int(datetime.now().timestamp())}")

def _gen_conn_id(user_id: int):
    return f'user-outland-{user_id}'

def _gen_user_client_cache_key(user_id: int):
    return f'user-outland-client-{user_id}'

@SingletonDecorator
class ConnectionManager:

  def __init__(self):
    self.channel = settings.WEBSOCKET_CHANNEL_NAME
    self.connections: Dict[str, Connection] = {}
    self.redis_manager: RedisManager | None = None

    thread = Thread(target=self.run_subscribe_in_thread)
    thread.start()

    ping_thread = Thread(target=self.run_ping_in_thread)
    ping_thread.start()

  async def connect(self, client_id:str,  websocket: WebSocket, user_id: int):
    if client_id in self.connections:
      raise Exception('client_id exist')
    await websocket.accept()
    conn = Connection(client_id, websocket)
    self.connections[client_id] = conn
    self.connections[_gen_conn_id(user_id)] = conn
    await RedisManager().set(_gen_user_client_cache_key(user_id), client_id)

  async def disconnect(self, client_id: str, user_id: int):
    if client_id in self.connections:
      try:
        conn = self.connections[client_id]
        await conn.websocket.close()
      except Exception as e:
        print(f"websocket disconnect, an unexpected error occurred: {e}")
        logger.error(f"IM websocket 断开出错: {e}")
      del self.connections[client_id]

      if user_id and _gen_conn_id(user_id) in self.connections:
        del self.connections[_gen_conn_id(user_id)]
        await RedisManager().delete(_gen_user_client_cache_key(user_id))

  def has_user_connection(self, user_id: int) -> bool:
    if _gen_conn_id(user_id) in self.connections:
      return True
    return False

  def has_connection(self, client_id: str, user_id: int | None) -> bool:
    if client_id in self.connections:
      return True
    if user_id:
      return self.has_user_connection(user_id)
    return False
  
  @staticmethod
  async def has_user_client(user_id: int) -> bool:
    return await RedisManager().exists(_gen_user_client_cache_key(user_id))
  
  @staticmethod
  async def get_user_client(user_id: int) -> str:
    return await RedisManager().get(_gen_user_client_cache_key(user_id))
  
  async def send_user_message(self, user_id: int, message: Message):
    conn_id = _gen_conn_id(user_id)
    if conn_id in self.connections:
      conn = self.connections[conn_id]
      await conn.send_message(message)

  async def send_message(self, message: Message):
    if message.client_id in self.connections:
      conn = self.connections[message.client_id]
      await conn.send_message(message)

  async def broadcast(self, message: Message):
    pass

  async def handle_message(self, message: Message):
    await self.send_message(message)

  async def publish(self, data):
    if self.redis_manager:
      await self.redis_manager.publish(self.channel, data)

  async def unsubscribe(self):
    await self.redis_manager.unsubscribe(self.channel)

  def run_subscribe_in_thread(self):
    if settings.REDIS_DB_ENABLE:
      self.redis_manager = RedisManager()
      asyncio.run(subscribe(self))

  def run_ping_in_thread(self):
    asyncio.run(ping(self))


async def ping(cm: ConnectionManager):
  while True:
    try:
      await asyncio.sleep(10)
      for key, value in cm.connections.items():
        if not key.startswith('user'):
          await value.ping()
    except Exception as e:
      print(f"An unexpected error occurred: {e}")


async def subscribe(cm: ConnectionManager):
  await cm.redis_manager.subscribe(cm.channel)

  while True:
    try:
      message = await cm.redis_manager.get_pubsub(cm.channel).get_message(ignore_subscribe_messages=True)
      if message:
        if message['type'] == 'message':
          data = Message.model_validate_json(message['data'])
          await cm.handle_message(data)
      else:
        await asyncio.sleep(0.1)  # 休眠以避免占用过多 CPU 资源
    except (ConnectionError, TimeoutError) as e:
      print(f"Connection error: {e}. Attempting to reconnect...")
      await cm.redis_manager.close()
      await cm.redis_manager.connect()
      await cm.redis_manager.subscribe(cm.channel)  # 重新订阅之前的频道
    except Exception as e:
      print(f"An unexpected error occurred: {e}")
      